package com.zxl.study.aops;

import com.zxl.study.annotations.RedisLimitAnnotation;
import com.zxl.study.exception.RedisLimitException;

import com.zxl.study.utils.RateLimitAlgorithm;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.DefaultParameterNameDiscoverer;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.expression.EvaluationContext;
import org.springframework.expression.Expression;
import org.springframework.expression.ExpressionParser;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.expression.spel.support.StandardEvaluationContext;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import org.springframework.core.io.ClassPathResource;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.lang.reflect.Method;
import java.util.*;

/**
 * @auther zxl
 * @create 2024-03-07 10:43
 */
@Aspect
@Component
public class RedisLimitAop {
    Object result = null;

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    private final Map<RateLimitAlgorithm, DefaultRedisScript<Long>> scripts = new HashMap<>();

    /**
     * 简单的令牌桶算法适合于实现快速限流，适用于请求速率稳定的简单场景。
     * 滑动窗口算法适合于需要精确控制请求速率的场景，特别是在严格时间窗口内的限流。
     * 漏桶算法适合于需要平滑突发流量的场景，可以防止突发流量对系统的冲击。
     */
    @PostConstruct
    public void init() {
        scripts.put(RateLimitAlgorithm.TOKEN_BUCKET, createRedisScript("rateLimiterTokenBucket.lua.lua"));
        scripts.put(RateLimitAlgorithm.SLIDING_WINDOW, createRedisScript("rateLimiterSlidingWindow.lua"));
        scripts.put(RateLimitAlgorithm.LEAKY_BUCKET, createRedisScript("rateLimiterLeakyBucket.lua"));
    }

    private DefaultRedisScript<Long> createRedisScript(String scriptFileName) {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(scriptFileName)));
        return redisScript;
    }

    @Around("@annotation(com.zxl.study.annotations.RedisLimitAnnotation)")
    public Object around(ProceedingJoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        //拿到RedisLimitAnnotation注解，如果存在则说明需要限流，容器捞鱼思想
        RedisLimitAnnotation redisLimitAnnotation = method.getAnnotation(RedisLimitAnnotation.class);

        if (redisLimitAnnotation != null) {
            //获取redis的key
            String key = redisLimitAnnotation.key();
            ExpressionParser parser = new SpelExpressionParser();

            Expression expression = parser.parseExpression(key);
            EvaluationContext context = new StandardEvaluationContext();
            Object[] args = joinPoint.getArgs();
            DefaultParameterNameDiscoverer discoverer = new DefaultParameterNameDiscoverer();
            String[] parameterNames = discoverer.getParameterNames(method);
            for (int i = 0; i < Objects.requireNonNull(parameterNames).length; i++) {
                context.setVariable(parameterNames[i], args[i].toString());
            }
            key = Objects.requireNonNull(expression.getValue(context)).toString();
            String className = method.getDeclaringClass().getName();
            String methodName = method.getName();
            String limitKey = key + "\t" + className + "\t" + methodName;

            if (null == key) throw new RedisLimitException("key cannot be null");

            long limit = redisLimitAnnotation.permitsPerSecond();
            long expire = redisLimitAnnotation.expire();
            RateLimitAlgorithm algorithm = redisLimitAnnotation.algorithm();

            List<String> keys = new ArrayList<>();
            keys.add(key);
            Long count;
            switch (algorithm) {
                case SLIDING_WINDOW:
                    count = stringRedisTemplate.execute(scripts.get(RateLimitAlgorithm.SLIDING_WINDOW), keys, String.valueOf(limit),
                            String.valueOf(System.currentTimeMillis()), String.valueOf(expire));
                    break;
                case LEAKY_BUCKET:
                    count = stringRedisTemplate.execute(scripts.get(RateLimitAlgorithm.LEAKY_BUCKET), keys, String.valueOf(limit),
                            String.valueOf(System.currentTimeMillis()), String.valueOf(expire));
                    break;
                case TOKEN_BUCKET:
                default:
                    count = stringRedisTemplate.execute(scripts.get(RateLimitAlgorithm.TOKEN_BUCKET), keys, String.valueOf(limit), String.valueOf(expire));
                    break;
            }

            if (count != null && count == 0) {
                return redisLimitAnnotation.msg();
            }
        }

        try {
            result = joinPoint.proceed();//放行
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }

        return result;
    }
}
